package com.example.harmtorch_demo.slice;

import com.example.harmtorch_demo.harmtorch.*;
import ohos.aafwk.ability.AbilitySlice;
import ohos.app.Context;
import ohos.global.resource.RawFileEntry;
import ohos.global.resource.Resource;
import ohos.hiviewdfx.HiLog;
import ohos.hiviewdfx.HiLogLabel;

import java.io.IOException;


public class Classifier {
    private Context context;
    private static final String MODEL_PATH = "entry/resources/rawfile/model.bin";
    private NN net = null;

    public Classifier(AbilitySlice slice) {
        this.context = slice.getContext();
    }

    public void init() {
        net = new NN(
                new Linear(28 * 28, 512),
                new Relu(),
                new Linear(512, 128),
                new Relu(),
                new Linear(128, 10),
                new Softmax()
        );
        try {
            loadModelFile();
        } catch (IOException e) {
            //e.printStackTrace();
        }

    }

    public Tensor getResult(Tensor input) {
        HiLogLabel label = new HiLogLabel(HiLog.LOG_APP, 0x00201, "test");
        // HiLog.warn(label, input.toString());
        Tensor output = net.inference(input);
        // HiLog.warn(label, output.toString());
        return output;
    }

    // Mappatura del file in memoria
    private void loadModelFile() throws IOException {
        RawFileEntry rawFileEntry = context.getResourceManager().getRawFileEntry(MODEL_PATH);
        try {
            Resource resource = rawFileEntry.openRawFile();
            net.readTorch(resource);
        } catch (Exception e) {
            //e.printStackTrace();
            System.exit(0);
        }
    }
}
